import torch.nn as nn
from mmdet.mmcv.utils import Registry
from mmdet.mmcv.cnn import MODELS as MMCV_MODELS
import mm_cls

MODELS = Registry('model', parent=MMCV_MODELS)

@MODELS.register_module()
class NetA(nn.Module):
    def forward(self, x):
        return x

net_a = MODELS.build(cfg=dict(type='NetA'))
print(net_a)
net_b = MODELS.build(cfg=dict(type='mm_cls.NetB'))
print(net_b)